import torch
from torch import Tensor, nn
import torch.nn.functional as F
from mmengine.model import BaseModule
from mmaction.registry import MODELS
from mmaction.utils import ConfigType
from typing import Optional, Sequence, Tuple, Union
from mmaction.utils import ForwardResults, SampleList

@MODELS.register_module()
class VenusFreqHead(BaseModule):
    """频率回归头，用于预测频率值
    
    Args:
        in_channels (int | Sequence[int]): 输入通道数
        loss (ConfigType): 损失函数配置，默认为 SmoothL1Loss
        hidden_channels (Sequence[int], optional): 隐藏层通道数
        dropout_ratio (float, optional): Dropout 比率
        init_cfg (ConfigType, optional): 初始化配置
    """

    def __init__(self,
                 in_channels: Union[int, Sequence[int]],
                 loss: ConfigType = dict(type='SmoothL1Loss'),
                 hidden_channels: Optional[Sequence[int]] = None,
                 dropout_ratio: float = 0.5,
                 init_cfg: Optional[ConfigType] = None):
        
        if init_cfg is None:
            init_cfg = self.default_init_cfg
            
        super().__init__(init_cfg=init_cfg)
        
        self.in_channels = in_channels
        self.hidden_channels = hidden_channels
        self.dropout_ratio = dropout_ratio
        
        # 构建损失函数
        self.loss_module = MODELS.build(loss)
        
        # 构建网络层
        self._build_layers()

    def _build_layers(self):
        """构建网络层"""
        layers = []
        
        # 如果输入是多尺度的特征，使用自适应平均池化将其统一到相同尺寸
        if isinstance(self.in_channels, Sequence):
            self.adaptive_pool = nn.AdaptiveAvgPool2d((1, 1))
            in_channels = sum(self.in_channels)
        else:
            self.adaptive_pool = None
            in_channels = self.in_channels
        
        # 添加隐藏层
        if self.hidden_channels:
            for hidden_dim in self.hidden_channels:
                layers.append(nn.Linear(in_channels, hidden_dim))
                layers.append(nn.ReLU(inplace=True))
                if self.dropout_ratio > 0:
                    layers.append(nn.Dropout(p=self.dropout_ratio))
                in_channels = hidden_dim
        
        # 添加输出层
        layers.append(nn.Linear(in_channels, 1))  # 输出一个频率值
        
        self.fc_layers = nn.Sequential(*layers)

    def forward(self, feats: Union[Tensor, Tuple[Tensor]], **kwargs) -> Tensor:
        """前向传播
        
        Args:
            feats (Tensor | Tuple[Tensor]): 主干网络提取的特征
            
        Returns:
            Tensor: 频率预测值，形状为 [batch_size, 1]
        """
        if isinstance(feats, Tuple):
            # 多尺度特征处理
            if self.adaptive_pool:
                # 对每个尺度的特征进行自适应平均池化并拼接
                pooled_feats = []
                for feat in feats:
                    pooled = self.adaptive_pool(feat)
                    pooled_feats.append(pooled)
                x = torch.cat(pooled_feats, dim=1)
            else:
                # 只使用最后一个尺度的特征
                x = feats[-1]
        else:
            # 单尺度特征
            x = feats
        
        # 展平特征
        x = torch.flatten(x, 1)
        
        # 通过全连接层
        freq_pred = self.fc_layers(x)
        
        return freq_pred

    def loss(self, feats: Union[torch.Tensor, Tuple[torch.Tensor]],
             data_samples: SampleList, **kwargs) -> dict:
        """计算损失
        
        Args:
            pred_freq (Tensor): 预测的频率值，形状为 [batch_size, 1]
            gt_freq (Tensor): 真实的频率值，形状为 [batch_size, 1]
            
        Returns:
            dict: 包含损失的字典
        """
        pred_freqs = self(feats, **kwargs)
        pred_freqs = pred_freqs.squeeze()

        gt_freqs = [x.gt_label[1] for x in data_samples]
        gt_freqs = torch.stack(gt_freqs).to(pred_freqs.device)
        gt_freqs = gt_freqs.squeeze()

        losses = dict()
        
        # 计算频率回归损失
        loss_freq = self.loss_module(pred_freqs, gt_freqs)
        losses['loss_freq'] = loss_freq
        
        return losses

    def predict(self, feats: Union[torch.Tensor, Tuple[torch.Tensor]],
                data_samples: SampleList, **kwargs) -> SampleList:
        """预测频率值
        
        Args:
            feats (Tensor | Tuple[Tensor]): 主干网络提取的特征
            
        Returns:
            Tensor: 频率预测值，形状为 [batch_size, 1]
        """
        pred_freqs =  self.forward(feats)

        for data_sample, pred_freq in zip(data_samples, pred_freqs):
            pred_freq = float(pred_freq)
            prev_label = data_sample.get('pred_label', None)
            prev_cls = int(prev_label[0]) if prev_label is not None else 0
            store_label = [prev_cls,pred_freq]
            data_sample.set_pred_label(store_label)
        return data_samples

    @property
    def default_init_cfg(self):
        """默认初始化配置"""
        init_cfg = [
            dict(type='Normal', layer='Linear', std=0.01, bias=0),
        ]
        return init_cfg